import sys
import subprocess
import platform

def check_python_version():
    """获取Python版本信息"""
    print("=== Python 版本信息 ===")
    print(f"Python 版本: {sys.version}")
    print(f"Python 解释器路径: {sys.executable}")
    print(f"操作系统: {platform.system()} {platform.release()}")
    print()

def check_cuda_version():
    """尝试获取CUDA版本信息"""
    print("=== CUDA 版本信息 ===")
    try:
        # 尝试通过nvidia-smi命令获取CUDA版本
        result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=True)
        output = result.stdout
        print("nvidia-smi 输出:")
        print(output)
        for line in output.split('\n'):
            if 'CUDA Version' in line:
                print(f"系统CUDA版本: {line.split(':')[-1].strip()}")
                break
    except (subprocess.SubprocessError, FileNotFoundError):
        print("未找到NVIDIA显卡或nvidia-smi命令不可用")
    
    # 尝试检查PyTorch中的CUDA版本
    try:
        import torch
        if torch.cuda.is_available():
            print(f"PyTorch CUDA 版本: {torch.version.cuda}")
            print(f"CUDA设备数量: {torch.cuda.device_count()}")
            for i in range(torch.cuda.device_count()):
                print(f"CUDA设备 {i}: {torch.cuda.get_device_name(i)}")
        else:
            print("PyTorch未启用CUDA支持")
    except ImportError:
        print("PyTorch未安装")
    print()

def check_pytorch_version():
    """获取PyTorch版本信息"""
    print("=== PyTorch 版本信息 ===")
    try:
        import torch
        print(f"PyTorch 版本: {torch.__version__}")
        
        # 尝试导入并获取torchvision版本
        try:
            import torchvision
            print(f"Torchvision 版本: {torchvision.__version__}")
        except ImportError:
            print("Torchvision 未安装")
        
        # 尝试导入并获取torchaudio版本
        try:
            import torchaudio
            print(f"Torchaudio 版本: {torchaudio.__version__}")
        except ImportError:
            print("Torchaudio 未安装")
        
        print(f"CUDA 是否可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"当前CUDA设备: {torch.cuda.current_device()}")
            print(f"当前CUDA设备名称: {torch.cuda.get_device_name(torch.cuda.current_device())}")
    except ImportError:
        print("PyTorch未安装")
    print()

def check_gpu_info():
    """获取显卡详细信息"""
    print("=== 显卡详细信息 ===")
    try:
        # 直接使用nvidia-smi命令获取显卡信息
        result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=True)
        print("显卡信息:")
        print(result.stdout)
        
        # 尝试获取更详细的信息，如显存使用情况
        detailed_result = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,memory.used,memory.free,utilization.gpu,utilization.memory", "--format=csv,noheader,nounits"], capture_output=True, text=True, check=True)
        print("\n详细显卡信息:")
        print("显卡名称,总显存(MB),已用显存(MB),可用显存(MB),GPU利用率(%),显存利用率(%)")
        print(detailed_result.stdout)
    except (subprocess.SubprocessError, FileNotFoundError):
        print("未找到NVIDIA显卡或nvidia-smi命令不可用")
        # 尝试使用dxdiag获取显卡信息（Windows系统）
        try:
            print("\n尝试通过dxdiag获取显卡信息...")
            # 输出到临时文件
            import tempfile
            import os
            temp_file = tempfile.NamedTemporaryFile(suffix='.txt', delete=False)
            temp_file.close()
            subprocess.run(["dxdiag", "/t", temp_file.name], capture_output=True, check=True)
            
            # 读取并显示显卡相关信息
            with open(temp_file.name, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                print("\n显卡信息摘要:")
                # 提取显示适配器部分
                adapter_section = False
                for line in content.split('\n'):
                    if 'Display Devices' in line:
                        adapter_section = True
                    elif adapter_section and line.strip() == '':
                        adapter_section = False
                    elif adapter_section:
                        print(line.strip())
            
            # 删除临时文件
            os.unlink(temp_file.name)
        except Exception as e:
            print(f"通过dxdiag获取显卡信息失败: {str(e)}")
    print()

def main():
    """主函数"""
    print("系统环境信息检测工具\n")
    check_python_version()
    check_pytorch_version()
    check_cuda_version()
    check_gpu_info()
    
    print("检测完成！")

if __name__ == "__main__":
    main()